#include "pathological.h"


#define MIN_ENODES 5
#define MAX_ENODES 50
#define MIN_LNODES 5

//#define SAMPLE_AGE 30

#define PIXEL_STRAT 3

float *lambdas;
camera_t camera;
double avglum = 1;
//int cachei, cachesize, tick;
//FILE *samplesf;
//double *lmap_cos;
int width, height;
color_t ambientior;
ALLEGRO_BITMAP *buffer;
ALLEGRO_DISPLAY *display;
refrac_stack_t *refracstack;
uint32_t refracstacksize = 10;
uint64_t passes = 0, paths = 0;
object_t **lights = NULL, **objects = NULL;
struct {uint32_t n; float X, Y, Z;} **image;
uint32_t numlights = 0, lightslots = 0, numobjects = 0, objectslots = 0;


void add_object(object_t *o)
{
	if(objectslots <= numobjects)
		objectslots = 1.2*(objectslots + 1),
		assert(objects = realloc(objects,objectslots*sizeof(object_t *)));
	*(objects + numobjects++) = o;

	if(o->mat.flags & MAT_EMITS)
	{
		if(lightslots <= numlights)
			lightslots = 1.2*(lightslots + 1),
			assert(lights = realloc(lights,lightslots*sizeof(object_t *)));
		*(lights + numlights++) = o;
	}
}

void push_refrac(object_t *o)
{
	int i;

	if(!refracstack->next) // Bigger!
	{
		refrac_stack_t *newstack
			= calloc(refracstacksize,sizeof(refrac_stack_t));
		for(i = 0; i < refracstacksize - 1; i++)
			(newstack + i)->next = newstack + i + 1,
			(newstack + i + 1)->prev = newstack + i;
		refracstack->next = newstack;
		newstack->prev = refracstack;
		refracstacksize *= 2;
	}

	(refracstack = refracstack->next)->o = o;
	refracstack->ior = &o->mat.ior;
}

void pop_refrac(object_t *o)
{
	refrac_stack_t *framep = refracstack;
	if(refracstack->o == o) refracstack = refracstack->prev;
	else
	{
		while(framep->o != o && framep->prev) framep = framep->prev;
		if(!framep->prev) return; // Don't remove the start of the stack
		framep->prev->next = framep->next;
		framep->next->prev = framep->prev;
		if(refracstack->next) refracstack->next->prev = framep;
		framep->next = refracstack->next;
		refracstack->next = framep;
		framep->prev = refracstack;
	}
}

int visible(point_t *p1, point_t *p2)
{
	ray_t r;
	double d;
	point_t p;
	object_t *o;
	r.o = *(vector_t *)p1;
	vec_sub((vector_t *)p2,(vector_t *)p1,&r.d);
	if(dot(&p1->n,&r.d) < 0 || dot(&p2->n,&r.d) > 0) return false;
	d = sqrt(dot(&r.d,&r.d));
	normalize_vec(&r.d);
	return !intersect_ray(&r,&p,&o) || d - p.t < EPSILON*p.t;
}

inline void use_sample(int x, int y, color_t *c)
{
	float X, Y, Z;
	color_to_XYZ(&X,&Y,&Z,c);
	if(x >= width) x = width - 1;
	if(y >= height) y = height - 1;
	image[x][y].X += X, image[x][y].Y += Y, image[x][y].Z += Z,
	image[x][y].n++;
}

// r->o.[xyz] == camera.[xyz]
inline void use_ray_sample(ray_t *r, color_t *c)
{
	int x, y;
	double mr1[4] = {camera.xstep.x,camera.ystep.x,-r->d.x,-camera.orgin.x},
	       mr2[4] = {camera.xstep.y,camera.ystep.y,-r->d.y,-camera.orgin.y},
	       mr3[4] = {camera.xstep.z,camera.ystep.z,-r->d.z,-camera.orgin.z};
	double *m[3] = {mr1,mr2,mr3};
	rref_3x4(m);
	if(0 <= (x = floor(m[0][3])) && x < width && 0 <= (y = floor(m[1][3]))
		&& y < height)
		use_sample(x,y,c);
}

// Random unit vector in cosine-weighted distribution in hemisphere around p.n
inline vector_t *direction_cos(vector_t *v, point_t *p)
{
	float r1 = ran(), r2 = ran();
	v->x = p->n.x + 2*cos(2*PI*r1)*sqrt(r2*(1 - r2));
	v->y = p->n.y + 2*sin(2*PI*r1)*sqrt(r2*(1 - r2));
	v->z = p->n.z + 1 - 2*r2;
	normalize_vec(v);
	return v;
}

inline vector_t *normalize_vec(vector_t *v)
{
	double s = 1/sqrt(sqr(v->x) + sqr(v->y) + sqr(v->z));
	v->x *= s, v->y *= s, v->z *= s;
	return v;
}

inline void sample_pixel(float x, float y)
{
	ray_t r;
	point_t p;
	object_t *o;
	vector_t id, od, v;
	color_t black, white;
	color_t ac, c, tc, bsdf;
	float atten, pdf, pcont;
	int d, i, n, ni, floorx = floor(x), floory = floor(y), side1, side2;
	struct {
		int id; // id is sequential from 0 at the eye/light
		point_t p;
		vector_t d; // From the previous node
		float pdf; // Of d
		object_t* o;
		color_t cb, cs; // Light reaching eye = cb + light leaving*cs
		color_t *ior; // Of -d
		int next;
		int prev;
	} nodes[MAX_ENODES];

	random_gray(&white,1); // Determines the wavelengths for the entire sample
	color_copy(&black,&white);
	color_scale(&black,0);
	color_copy(&bsdf,&white);
	lambdas = white.w;

	while(refracstack->prev) refracstack = refracstack->prev; 

	// Trace the eye path, marking each hit along the way
	r.o.x = camera.x, r.o.y = camera.y, r.o.z = camera.z;
	r.d.x = camera.orgin.x + x*camera.xstep.x + y*camera.ystep.x;
	r.d.y = camera.orgin.y + x*camera.xstep.y + y*camera.ystep.y;
	r.d.z = camera.orgin.z + x*camera.xstep.z + y*camera.ystep.z;
	normalize_vec(&r.d);
	nodes[0].id = 0, nodes[0].o = NULL;
	nodes[0].p.x = camera.x, nodes[0].p.y = camera.y, nodes[0].p.z = camera.z;
	nodes[0].p.n.x = 0, nodes[0].p.n.y = 0, nodes[0].p.n.z = 0;
	color_copy(&nodes[0].cb,&black);
	color_copy(&nodes[0].cs,&white);
	color_mul(&nodes[0].cs,&ambientior), color_mul(&nodes[0].cs,&ambientior);
	nodes[0].ior = &ambientior, nodes[0].prev = -1, pdf = 1;
	for(ni = n = 1; n < MAX_ENODES && (n < MIN_ENODES || ran() < 0.5);
		ni++, n++)
	{
		if(!intersect_ray(&r,&nodes[ni].p,&nodes[ni].o)) break;
		nodes[ni].id = n, nodes[ni].d = r.d, nodes[ni].pdf = pdf,
		nodes[ni-1].next = ni, nodes[ni].prev = ni - 1;

		// Migrate the importance
		pcont = n < MIN_ENODES ? 1 : 0.5;
		if(nodes[ni-1].o)
			color_copy(&nodes[ni].cb,nodes[ni-1].ior),
			color_mul(&nodes[ni].cb,&nodes[ni].cb),
			color_rcp(&nodes[ni].cb),
			color_mul(&nodes[ni].cb,&nodes[ni-1].o->mat.emission),
			color_mul(&nodes[ni].cb,&nodes[ni-1].cs),
			color_add(&nodes[ni].cb,&nodes[ni-1].cb),
			atten = fabs(dot(&nodes[ni-1].p.n,&nodes[ni].d)),
			nodes[ni-1].o->mat.bsdf(&bsdf,&nodes[ni-1].o->mat,&nodes[ni-1].p,
				&nodes[ni-1].d,&nodes[ni].d),
			refracstack->o // Beer's law
				? color_copy(&nodes[ni].cs,&refracstack->o->mat.transmit),
					color_scale(&nodes[ni].cs,-nodes[ni-1].p.t),
					color_exp(&nodes[ni].cs),
					color_mul(&nodes[ni].cs,&nodes[ni-1].cs)
				: color_copy(&nodes[ni].cs,&nodes[ni-1].cs),
			color_mul(&nodes[ni].cs,&bsdf),
			color_scale(&nodes[ni].cs,atten/pdf/pcont),
			// Deal with the refrac stack
			side1 = dot(&nodes[ni-1].d,&nodes[ni-1].p.n) < 0,
			side2 = dot(&nodes[ni].d,&nodes[ni-1].p.n) > 0,
			side1 && !side2 ? push_refrac(nodes[ni].o)
				: !side1 && side2 ? pop_refrac(nodes[ni].o) : 0;
		else color_copy(&nodes[ni].cb,&nodes[ni-1].cb),
			color_copy(&nodes[ni].cs,&nodes[ni-1].cs),
			color_scale(&nodes[ni].cs,1/pcont);

		nodes[ni].ior = refracstack->ior;

		// Figure out where to go from here
		id = r.d;
		r.o = *(vector_t *)&nodes[ni].p;
		pdf = nodes[ni].o->mat.sample_bsdf(&r.d,&nodes[ni].o->mat,&nodes[ni].p,
			&id);
	}
	nodes[ni-1].next = -1;

	while(refracstack->prev) refracstack = refracstack->prev;
	(i = floor(ran()*(numlights))) >= numlights ? (i = numlights - 1) : 0;
	o = *(lights + i);
	o->surface_point(&p,o);
	direction_cos(&id,&p);
	id.x = -id.x, id.y = -id.y, id.z = -id.z;
	color_copy(&c,&black);
	pdf = 1;
	for(d = 0; d < MIN_LNODES || ran() < 0.5; d++)
	{
		// Connect this point to each eye-path node
		r.o.x = camera.x, r.o.y = camera.y, r.o.z = camera.z;
		pcont = d < MIN_LNODES ? 1 : 0.5;
		for(ni = 0; ni >= 0; ni = nodes[ni].next)
		{
			if(visible(&p,&nodes[ni].p))
			{
				atten = fabs(dot(&p.n,&id));
				normalize_vec(vec_sub((vector_t *)&nodes[ni].p,(vector_t *)&p,
					&od));
				o->mat.bsdf(&bsdf,&o->mat,&p,&id,&od);
				color_copy(&ac,&c);
				color_mul(&ac,&bsdf);
				color_scale(&ac,atten/pdf);
				color_add(&ac,&o->mat.emission);
				color_scale(&ac,1/pcont);

				// Beer's law
				if(refracstack->o)
					color_copy(&tc,&refracstack->o->mat.transmit),
					color_scale(&tc,-p.t),
					color_exp(&tc),
					color_mul(&ac,&tc);

				if(nodes[ni].id)
					atten = fabs(dot(&nodes[ni].p.n,&od)),
					v.x = -nodes[ni].d.x,
					v.y = -nodes[ni].d.y,
					v.z = -nodes[ni].d.z,
					nodes[ni].o->mat.bsdf(&bsdf,&nodes[ni].o->mat,&nodes[ni].p,
						&od,&v),
					color_mul(&ac,&bsdf),
					color_scale(&ac,atten/nodes[ni].pdf),
					color_add(&ac,&nodes[ni].o->mat.emission),
					color_copy(&tc,nodes[ni].ior),
					color_mul(&tc,&tc), color_rcp(&tc), color_mul(&ac,&tc),
					color_mul(&ac,&nodes[ni].cs),
					color_add(&ac,&nodes[ni].cb),
					use_sample(floorx,floory,&ac);
				else vec_sub((vector_t *)&p,&r.o,&r.d), use_ray_sample(&r,&ac);
			}
			else if(nodes[ni].id)
				color_copy(&ac,nodes[ni].ior),
				color_mul(&ac,&ac),
				color_rcp(&ac),
				color_mul(&ac,&nodes[ni].o->mat.emission),
				color_mul(&ac,&nodes[ni].cs),
				color_add(&ac,&nodes[ni].cb),
				use_sample(floorx,floory,&ac);

		}

		r.o = *(vector_t *)&p;
		pdf = o->mat.sample_bsdf(&r.d,&o->mat,&p,&id);
		// Migrate the color through the scene
		atten = fabs(dot(&p.n,&id));
		o->mat.bsdf(&bsdf,&o->mat,&p,&id,&r.d);
		color_mul(&c,&bsdf),
		color_scale(&c,atten/pdf),
		color_add(&c,&o->mat.emission),
		color_scale(&c,1/pcont);
		// Deal with the refrac stack
		side1 = dot(&id,&p.n) < 0;
		side2 = dot(&r.d,&p.n) > 0;
		if(side1 && !side2) push_refrac(o);
		else if(!side1 && side2) pop_refrac(o);
		// Check for the next point
		if(!intersect_ray(&r,&p,&o)) break;
		// Beer's law
		if(refracstack->o)
			color_copy(&tc,&refracstack->o->mat.transmit),
			color_scale(&tc,-p.t),
			color_exp(&tc),
			color_mul(&c,&tc);
		id = r.d;

	}

	paths++;
}

void *render_thread(ALLEGRO_THREAD *thread, void *arg)
{
	time_t t = time(NULL);
	int x = 0, y = 0, i, j;

	puts("Rendering...");

	image = malloc(width*sizeof(void *));
	for(i = 0; i < width; i++)
		image[i] = calloc(sizeof(**image),height);

	refracstack = calloc(refracstacksize,sizeof(*refracstack));
	for(i = 0; i < refracstacksize - 1; i++)
		(refracstack + i)->next = refracstack + i + 1,
		(refracstack + i + 1)->prev = refracstack + i;
	refracstack->ior = &ambientior;

	while(!al_get_thread_should_stop(thread))// && difftime(time(NULL),t) < 10)
	{
		for(i = 0; i < PIXEL_STRAT; i++)
			for(j = 0; j < PIXEL_STRAT; j++)
				sample_pixel(x + (i + ran())/PIXEL_STRAT,
					y + (j + ran())/PIXEL_STRAT);

		if(++x >= width) x = 0, y++;
		if(y >= height) y = 0, passes++;
	}
	printf("Rendering finished in %g seconds\n",difftime(time(NULL),t));

	return NULL;
}

void update_buffer()
{
	int x, y;
	double lum = 0;
	struct {float r, g, b;} c;
	al_set_target_bitmap(buffer);
	for(y = 0; y < height; y++)
		for(x = 0; x < width; x++)
		{
			if(!image[x][y].n || (image[x][y].X == 0
				&& image[x][y].Y == 0 && image[x][y].Z == 0))
				{al_put_pixel(x,y,al_map_rgb(0,0,0)); continue;}
			uint32_t n = image[x][y].n;
			float X = image[x][y].X/n, Y = image[x][y].Y/n,
				Z = image[x][y].Z/n;
			XYZ_to_rgb(&c.r,&c.g,&c.b,&X,&Y,&Z);
			if(c.r >= 0 && c.r < DBL_MAX
				&& c.g >= 0 && c.g < DBL_MAX
				&& c.b >= 0 && c.b < DBL_MAX)
			{
				lum += 0.2126*c.r + 0.7152*c.g + 0.0722*c.b;
			}//else raise(SIGINT);
			if((c.r /= 2*avglum) > 1) c.r = 1;
			if((c.g /= 2*avglum) > 1) c.g = 1;
			if((c.b /= 2*avglum) > 1) c.b = 1;

			//if(!n) al_put_pixel(x,y,al_map_rgb(0,0,0));
			//else al_put_pixel(x,y,al_map_rgb(255,255,255));
			al_put_pixel(x,y,al_map_rgb_f(c.r,c.g,c.b));
		}
	avglum = lum/(width*height);
	al_set_target_backbuffer(display);
	al_draw_bitmap(buffer,0,0,0);
	al_flip_display();

	printf("Refreshed display; %.f pass%s completed, %.f paths traced.\n",
		(double)passes,passes == 1 ? "":"es",(double)paths);
}

int main(int argc, char **argv)
{
	char *ext;
	int t, x, y;
	int quit = false;
	FILE *out, *scene;
	ALLEGRO_EVENT_QUEUE *queue;
	ALLEGRO_THREAD *renderthread;

	width = 480;
	height = 360;

	ran_init();

	cieXYZ_init();

	random_gray(&ambientior,1);

	// Load the model
	printf("Parsing input file...");
	assert(argc >= 2);
	assert(scene = fopen(argv[1],"r"));
	assert(ext = strrchr(argv[1],'.'));
	if(strcmp(ext,".msdl") == 0) msdl(scene), msdl_init();
	else if(strcmp(ext,".pwld") == 0) pwld_import(scene);
	else printf("error: unknown file extension '%s'",ext), exit(1);
	fclose(scene);
	printf("done.\n");

	bih_build();

	// Set up Allegro stuff
	assert(al_init());
	assert(al_install_mouse());
	assert(al_install_keyboard());

	// Set up the display
	al_set_new_display_flags(ALLEGRO_WINDOWED);
	assert(display = al_create_display(width,height));

	// Set up the rendering thread
	t = al_get_new_bitmap_flags();
	al_set_new_bitmap_flags(ALLEGRO_MEMORY_BITMAP);
	assert(buffer = al_create_bitmap(width,height));
	al_set_new_bitmap_flags(t);
	renderthread = al_create_thread(&render_thread,NULL);
	al_start_thread(renderthread);

	// Set up the event queue
	assert(queue = al_create_event_queue());
	al_register_event_source(queue,al_get_keyboard_event_source());

	// Do the event loop
	while(!quit)
	{
		ALLEGRO_EVENT e;
		al_wait_for_event(queue,&e);
		switch(e.type)
		{
			case ALLEGRO_EVENT_KEY_CHAR:
				switch(e.keyboard.keycode)
				{
					case ALLEGRO_KEY_ESCAPE:
						quit = true;
						break;
					case ALLEGRO_KEY_ENTER:
						update_buffer();
						break;
				}
				break;
		}
	}

	// Bye... :(
	al_destroy_thread(renderthread);
	update_buffer();
	assert(out = popen("pfsout out.hdr","w"));
	fprintf(out,"PFS1\n%u %u\n3\n2\nLUMINANCE=RELATIVE\nWHITE.Y=1\nX\n0\nY\n0"
		"\nZ\n0\nENDH",width,height);
	for(y = 0; y < height; y++)
		for(x = 0; x < width; x++)
		{
			float X = image[x][y].X/image[x][y].n;
			assert(fwrite(&X,4,1,out) == 1);
		}
	for(y = 0; y < height; y++)
		for(x = 0; x < width; x++)
		{
			float Y = image[x][y].Y/image[x][y].n;
			assert(fwrite(&Y,4,1,out) == 1);
		}
	for(y = 0; y < height; y++)
		for(x = 0; x < width; x++)
		{
			float Z = image[x][y].Z/image[x][y].n;
			assert(fwrite(&Z,4,1,out) == 1);
		}
	pclose(out);
	printf("Image written to out.hdr\n");
	al_destroy_bitmap(buffer);

//	fclose(samplesf);

	return 0;
}

